#!/usr/bin/env python
"""
SDRF tools

 Tools to assist in SDRF parsing

  Currently supports translation from TCGA UUID to barcodes and SDRF
  crunching and lookup

Author: Tom Robinson <trobinson@systemsbiology.org>
Initial Revision: January 13th, 2012

TODO:
* Support cloning, recreation, ___repr__, ___str__, et cetera
* Support threading

"""

## Imports

import json
import logging
import os
import re
import urllib2


## Credit documentation

__author__ = "Tom Robinson"
__copyright__ = "Institute for Systems Biology, 2012"
__credits__ = ["Tom Robinson"]
__license__ = "Apache 2.0"
__version__ = "1.0"
__maintainer__ = "Tom Robinson"
__email__ = "trobinson@systemsbiology.org"
__status__ = "Production"
__platform__ = "Multi-OS, Python 2.6+"


## Constants

LINE_SEPARATOR = "\n"           ## Line separator: UNIX-style line breaks
LOOKUP_SERVICE_URI = 'https://tcga-data.nci.nih.gov/uuid/uuidws/mapping/' \
                      'json/uuid/'
                                ## URI for UUID-to-barcode translation requests
MAX_LINE=40960                  ## Maximum line length (in characters)

BARCODE_RE = r"tcga-[0-9a-z]{2}-[0-9a-z]{4}-[0-9a-z]{2}[0-9A-Z-]*"
                                ## TCGA barcode regular expression, at least
                                ##  sample length
UUID_RE = r"[0-9abcdef]{8}-[0-9abcdef]{4}-[0-9abcdef]{4}-" \
            + "[0-9abcdef]{4}-[0-9abcdef]{12}"
                                ## UUID regular expression

COMBINED_LOOKUP_RE = r"(" + BARCODE_RE + r"|" + UUID_RE + r")"
                                ## Combined UID regular expression


## Classes

class Barcodizer:
    """ Converts UUIDs to barcodes, by cache and web lookup

        This abstraction allows for the use of a shared object that controls a
        cache file and web lookup service with adherence to RAII. In this way,
        we can pass this object around when we wish to perform cached UUID to
        barcode translation without forcing the caller to have to manage the
        cache itself.

    """

    _logger = None          ## Logger to write any informative messages into 
                            ##  (see __init__) ##
    _cache = {}             ## Internal cache used for UUID-barcode translation
                            ##  (see __init__) ##
    _cacheFile = None       ## Cache filepath for UUID-barcode translation
                            ##  (see __init__) ##
    _cacheFHandle = None    ## Cache file handle to write new lookups into
                            ##  (see __init__) ##

    def __init__(self, cacheFile=None, logger=None):      
        """ Create a new barcodizer

        Parameters
        ----------
        cacheFile : string
            Cache filepath to read from and write into. If this is left as
            None, no cache will be used (all lookups will be done live).
            Default: None

        logger : string
            Logger (from module "logging") passed in from the caller.
            Default: None

        Return Value
        ------------
        SdrfIndexer
            A new instance of this class

        Examples
        --------
        >>> b = Barcodizer()
        
        """

        # Set our internal variables
        #  Currently, incorrect typing is the caller's problem
        self._logger = logger
        self._cacheFile = cacheFile

        # Propagate our internal cache if the file already exists on disk
        if self._cacheFile and os.path.isfile(self._cacheFile):

            # Open the read handle
            if self._logger:
                self._logger.info( \
                  'Found cache file at %s; reading its contents' \
                  % self._cacheFile)
            fhandle = open(self._cacheFile,'rU') 

            # Read in the file
            line = '\n'
            while line:
                line = fhandle.readline(MAX_LINE)
                if not line:
                    break
                if len(line) == MAX_LINE:
                    raise IOError('Line length was too large (%s).' \
                      % str(MAX_LINE) \
                      + 'Please ensure the file is being parsed correctly.')

                # Ensure the file isn't corrupt
                split = line.split('\t')
                if len(split) != 2:            
                    raise IOError('Cache file on disk is corrupt.' \
                      + ' Please fix or remove it before continuing.')

                # Add this result to our cache
                self._cache[split[0]] = split[1].rstrip()
        
            # We're done reading; close the read handle
            fhandle.close()
            if self._logger:
                self._logger.info('Done reading %s' \
                  % self._cacheFile)

            # ... and open our internally-pointed append handle
            if self._logger:
                self._logger.info('Setting %s to append now' \
                  % self._cacheFile)
            self._cacheFHandle = open(self._cacheFile,'a')

    
    def extractRawIdentifiersFromString(self, data):
        """ Extract raw barcodes or UUIDs from the given string

        Parameters
        ----------
        data : string
            String of data to perform identifier search within

        Return Value
        ------------
        list of string
            All identifiers that could be extracted from the string

        Examples
        --------
        >>> myString = "TCGA-00-0000-00"
        >>> b = Barcodizer()
        >>> b.extractRawIdentifiersFromString(myString) # ids are extracted
        ['tcga-00-0000-00']

        """
        # Match TCGA barcodes and UUIDs found in this line
        matches = re.findall(COMBINED_LOOKUP_RE, data.lower())

        if not matches:
            return []

        return list(matches)

    
    def extractBarcodesFromString(self, data):
        """ Extract the unique TCGA barcodes from a string of data

         If the given value isn't a barcode, this will attempt to look it up
         using the TCGA UUID service. If that fails, an IOError is thrown.

        Parameters
        ----------
        data : string
            String of data to perform barcode search within

        Return Value
        ------------
        list of string
            All barcodes that could be extracted from the string

        Examples
        --------
        >>> myString = "TCGA-00-0000-00"
        >>> b = Barcodizer()
        >>> b.extractBarcodesFromString(myString)   # barcodes are extracted 
        ['tcga-00-0000-00']
        
        """
        # Wrap extractRawIdentifiersFromString(...), providing the additional
        #  translation calls necessary to turn this into a barcode
        ret = set()
        for identifier in self.extractRawIdentifiersFromString(data):

            # Look it up, if it happens to be a UUID
            if len(identifier) == 36:
                try:
                    identifier = self.translate(identifier)
                except IOError:
                    if self._logger:
                        self._logger.warn( \
                          "Translation of UUID failed during web lookup!" \
                          + " UUID %s will be omitted from processing." % \
                          identifier)
                    continue

            ret.add(identifier)

        return list(ret)


    def translate(self, uuid):
        """ Perform UUID to TCGA barcode translation

         This method will only call out to a web service if a cache miss occurs
         or no cache file was supplied. New lookups will be cached in memory
         and, when supplied, to the cache file as well. 

         Please ensure locked access to your cache file before running
         translate, or file corruption is exceedingly likely.

         Notes:

         * THIS WILL BREAK: the UUID-to-barcode lookup service is
          unstable and incomplete! We use this to translate to TCGA
          barcodes, but there are no guarantees on further support. (!!)
        
         * We assume that if the cache file breaks, we should simply
          re-baseline from the web service. Thus, no timestamp information
          is preserved.

        Parameters
        ----------
        uuid : string
            UUID to translate into a TCGA barcode

        Return Value
        ------------
        string
            The TCGA barcode identified for this UUID

        Examples
        --------
        >>> myuuid = "ff4890d1-67e8-4d62-b116-2b049753f6ee"
        >>> b = Barcodizer()
        >>> b.translate(myuuid)     # myuuid is translated into a barcode
        u'tcga-28-1746-01a-21-1898-20'

        """
        
        if not (len(uuid) == 36 and re.match(UUID_RE,uuid)):
            raise ValueError("Supplied UUID %s was invalid" % uuid)
        else:
            if uuid not in self._cache:
                if self._logger:
                    self._logger.info( \
                        'UUID %s not in cache; retrieving it, then.' % \
                        uuid) 

                barcode = None
                try:                    
                    conn = urllib2.urlopen(LOOKUP_SERVICE_URI + uuid) 
                    barcode = \
                        unicode(json.loads(conn.read())['barcode'].lower())
                except: 
                    raise IOError( \
                        "Could not complete translation request for UUID" \
                            + " %s" % uuid)
                    
                # If our cache file is specified, write into it
                if self._cacheFHandle:
                    self._cacheFHandle.write(uuid)
                    self._cacheFHandle.write('\t')
                    self._cacheFHandle.write(barcode)
                    self._cacheFHandle.write(LINE_SEPARATOR)

                # Add it to our internal cache, as well
                self._cache[uuid] = barcode

                # Reassign our unique identifier
                return barcode

            else:
                # Just reassign our unique identifier
                return self._cache[uuid]


    def __del__(self):
        """ Object destructor

         Ensures our cache file is closed after writing, if one was specified.

        Parameters
        ----------
        None (other than the built-in, "self")

        Return Value
        ------------
        None (this object is destroyed)

        Examples
        --------
        >>> b = Barcodizer()
        >>> del b   # this Barcodizer is now destroyed

        """

        # If a cache file was being written to, close it
        if self._cacheFHandle:
            self._cacheFHandle.close()


class SdrfIndexer:
    """ An SDRF file reader that maintains a dictionary of archives and
     filenames found in SDRF files

     This class is useful primarily as an indexing service for filepaths that
     must have their unique identifiers defined in an SDRF file or collection
     of SDRF files. See the remainder of the class API for more information. 

     This class is designed to only output a singular barcode match per
     file. If multiple matches are detected for a single filepath, you'll
     receive a warning and None on lookup, because in those cases the
     identification is vacuous outside of debugging and consistency checks.

     Note: this class makes no guarantees about thread safety and only minimal
     guarantees about performance. Caveat emptor. (!!)

    """
    _archiveDict = {}  ## Dictionary of archive directories, containing a list
                        ##  of their associated identifiers ##
    _filenameDict = {} ## Dictionary of exact filenames, containing a list of
                        ##  their associated identifiers ##

    _logger = None     ## Logger to write any informative messages into
                        ##  (see __init__) ##

    def __init__(self, barcodizer, logger=None):
        """ Create a new indexer 

        Parameters
        ----------
        barcodizer : Barcodizer
            Translation service to identify TCGA barcodes from UUIDs

        logger : string
            Logger (from module "logging") passed in from the caller
            Default: None

        Return Value
        ------------
        SdrfIndexer
            A new instance of this class

        Examples
        --------
        >>> si = SdrfIndexer(Barcodizer())
        
        """
        self._logger = logger
        self._barcodizer = barcodizer


    def appendFromFile(self, filepath):
        """ Add the contents of an SDRF file to our internal dictionaries 

        Parameters
        ----------
        filepath : string
            Path to the target file to open for reading

        Return Value
        ------------
        None (contents are appended if an exception isn't raised)

        Examples
        --------
        >>> si = SdrfIndexer(Barcodizer())
        >>> si.appendFromFile('/dev/null')

        """ 
        if self._logger: 
            self._logger.info('appendFromFile: attempting to read %s' % \
                filepath) 

        fhandle = open(filepath,'rU')
        self.appendFromStream(fhandle)
        fhandle.close()
        
        if self._logger: 
            self._logger.info('appendFromFile: finished reading %s' % \
                filepath) 


    def appendFromStream(self, stream, lines=-1):
        """ Add the contents of an SDRF stream to the internal dictionaries 

        Parameters
        ----------
        stream : streamlike
            An open, readable stream object

        lines : int
            Number of lines to read from the stream. Default: -1 (disable)

        Return Value
        ------------
        None (contents are appended if an exception isn't raised)

        Examples
        --------
        >>> si = SdrfIndexer(Barcodizer())
        >>> mystream = open('/dev/null','rU')
        >>> si.appendFromStream(mystream)
        >>> mystream.close()

        """ 
        if not isinstance(lines,int) and lines < -1:
            raise ValueError("Number of lines specified was invalid: should be"
                + " positive, or -1 to disable.")

        header = stream.readline(MAX_LINE)
        if len(header) == MAX_LINE:
            raise IOError('Line length was too large (%s).' \
              % str(MAX_LINE) \
              + 'Please ensure the file is being parsed correctly.')
        parsedHeader = header.lower().split('\t')

        tcgaArchiveIdxs = []
        tcgaFilepathIdxs = []
        for idx, columnHeading in enumerate(parsedHeader):
            # These headings are given here:
            #  https://wiki.nci.nih.gov/display/TCGA/Sample+and+Data+Relationship+Format 
            if '[tcga archive name]' in columnHeading:
                tcgaArchiveIdxs.append(idx)
            if 'file' in columnHeading:
                tcgaFilepathIdxs.append(idx)

        line = '\n'
        lineCount = 0
        while line:

            # Honor the number of lines if this feature is enabled
            if lines != -1:
                if lines == 0:
                    break
                lines -= 1

            # Read a line; increment the line count
            line = stream.readline(MAX_LINE)
            if not line:
                break
            if len(line) == MAX_LINE:
                raise IOError('Line length was too large (%s).' \
                  % str(MAX_LINE) \
                  + 'Please ensure the file is being parsed correctly.')
       
            lineCount += 1         

            # Match TCGA barcodes and UUIDs found in this line
            found_identifiers = \
              self._barcodizer.extractBarcodesFromString(line) 

            if not found_identifiers:
                #print line
                if self._logger:
                    self._logger.warn( \
                        'TCGA barcode or UUID not found on line'
                        + ' ' + str(lineCount) + '.'
                        +  ' Ignoring line and continuing happily...')
                continue

            if len(found_identifiers) > 1:
                if self._logger:
                    self._logger.warn( \
                        'Inconsistent variety of barcodes and UUIDs'
                        + ' found on line ' + str(lineCount) + '.'
                        + ' Ignoring line and continuing happily...')
                continue    

            # Re-address the now-identified unique identifier
            unique_identifier = found_identifiers[0]

            columnData = line.rstrip().split('\t')

            # Walk the previously-identified indices, searching for archive data
            for idx in tcgaArchiveIdxs:
                if idx >= len(columnData):
                    if self._logger:
                        self._logger.warn( \
                            'Header length exceeded length of line'
                            + ' ' + str(lineCount) + '!'
                            + ' Ignoring line and continuing happily...')
                    break

                # Sentinel value explicitly defined as containing no
                #  information in MAGE-TAB files. If we find it, continue.
                if columnData[idx] == '->':
                    continue

                if self._logger:
                    self._logger.info('Appending %s to _archiveDict[%s]' % \
                        (unique_identifier,columnData[idx]))
                if columnData[idx] not in self._archiveDict:
                    self._archiveDict[columnData[idx]] \
                         = [unique_identifier]
                else:
                    self._archiveDict[columnData[idx]].append( \
                        unique_identifier)


            # Walk the previously-identified indices, searching for file data
            for idx in tcgaFilepathIdxs:
                if idx >= len(columnData):
                    if self._logger:
                        self._logger.warn( \
                            'Header length exceeded length of line'
                            + ' ' + str(lineCount) + '!'
                            + ' Ignoring line and continuing happily...')
                    break
    
                # Sentinel value explicitly defined as containing no
                #  information in MAGE-TAB files. If we find it, continue.
                if columnData[idx] == '->':
                    continue

                if self._logger:
                    self._logger.info('Appending %s to _filenameDict[%s]' % \
                        (unique_identifier,columnData[idx]))
                if columnData[idx] not in self._filenameDict:
                    self._filenameDict[columnData[idx]] \
                         = [unique_identifier]
                else:
                    self._filenameDict[columnData[idx]].append( \
                        unique_identifier)

        
    def lookup(self, filepath):
        """ Attempt to look up the given file path in our internal dictionaries

         If the lookup fails, None will be returned. If a logger is supplied,
         details of the lookup failure will be written to the log.

         NOTE ALSO: filepath lookups are case-sensitive.

        Parameters
        ----------
        filepath : string
            The given filepath to look up in our internal dictionaries

        Return Value
        ------------
        string
            If the path can be uniquely identified, its identifier will be
            returned. Otherwise, None.

        Examples
        --------
        >>> si = SdrfIndexer(Barcodizer())
        >>> si.lookup('/my/TCGA/files/here')

        """
        uniqueIdentifier = None        
        by_archive = None
        by_filename = None
        exploded_path = filepath.split(os.sep)

        if len(exploded_path) == 0:
            raise ValueError("Given path %s was invalid" % filepath)

        for pathpiece in exploded_path:

            # Ignore vacuous path pieces
            if not pathpiece:
                continue

            # Look up this piece in our dictionaries
            if pathpiece in self._archiveDict:
                if by_archive:
                    if self._logger:
                        self._logger.warn( \
                            'Found multiple matches for this archive.'
                            + ' Please inspect the input SDRF file source list'
                            + ' for %s. Continuing happily...' % filepath)
                    return None

                if self._logger:
                    self._logger.info('Found %s in self._archiveDict' % \
                        pathpiece) 
                by_archive = dict([[entry,''] for entry \
                    in self._archiveDict[pathpiece]])

        # Additionally, see if the filename matches
        if exploded_path[-1] and exploded_path[-1] in self._filenameDict:
            if self._logger:
                self._logger.info('Found %s in self._filenameDict' % \
                    pathpiece)
            by_filename = self._filenameDict[pathpiece]
        
        # Perform dict indexing against the archive collection only, since this
        #  is more likely to be the larger collection
        if by_archive and by_filename:
            for id in by_filename:
                if id in by_archive:
                    if uniqueIdentifier and uniqueIdentifier != id:
                        if self._logger:
                            self._logger.warn( \
                                'Multiple irreconcilable identifiers found.'
                                + ' This file may be a composite. Please inspect'
                                + ' the input SDRF file source list for' 
                                + ' %s. Continuing happily...' % \
                                filepath) 
                        return None
                    uniqueIdentifier = id

        if self._logger:
            if uniqueIdentifier:
                self._logger.info( \
                    'Successfully looked up this file as %s' % uniqueIdentifier)
            else:
                self._logger.info( \
                    'Unable to identify this file via SDRF parsing')

        return uniqueIdentifier


# Invoke our inline doctests when this library is invoked directly
if __name__ == "__main__":
    import doctest
    doctest.testmod()
